
import os
import json
import torch
import numpy as np
from torch.utils.data import Dataset
from PIL import Image


class PieceData:

    def __init__(self,data:dict,data_dir:str,query_template:str) -> None:
        self.data_dir = data_dir
        self.data = data
        self.image_name = data["image"]
        self.image_path = os.path.join(self.data_dir,self.image_name)
        self.answer_str = data["answer"]
        self.tags = set(data.keys()) - set(["image","answer"])
        self.tag_data = { key : val for key,val in data.items() if key in self.tags }
        self.query_template = query_template
        self._assign_tag_data_to_query_template()
        pass

    def _assign_tag_data_to_query_template(self):
        query = str(self.query_template)
        for key,val in self.tag_data.items():
            key = f"<{key}>"
            if key in query:
                query = query.replace(key,val,-1)
            else:
                print(f"key {key} is not in query template: {self.query_template}")
        self.query = query
    
    def get_dict(self,trans=None) -> dict:
        image = Image.open(self.image_path).convert("RGB")
        if trans is not None:
            image = trans(image)
        return {
            "query" : self.query,
            "answer": self.answer_str,
            "image_name": self.image_name,
            "image_path": self.image_path,
            "image": image
        }


class PatchDataSet(Dataset):
    def __init__(self, query_path, data_path, trans=None):
        self.query_path = query_path
        self.data_path = data_path
        self.trans = trans
        self.all_data:list[PieceData] = []
        self.length = 0
        self._load()

    def _load(self):
        content = None
        with open(self.query_path, "r", encoding="utf-8") as fp:
            content = json.load(fp)
            content = dict(content)
        for dirname,data in content.items():
            query_template = str(data["query_template"])
            for piece in list(data["data"]):
                self.all_data.append(PieceData(
                    piece,
                    os.path.join(self.data_path,dirname),
                    query_template
                ))
        self.length = len(self.all_data)

    def __len__(self):
        return len(self.all_data)

    def __getitem__(self, index):
        item = dict()
        if index < self.length and index >= 0:
            item = self.all_data[index].get_dict(self.trans)
        return item


pds = PatchDataSet(
    "./query.json",
    "./",
    None
)

print(len(pds))
for idx in range(len(pds)):
    print(pds.__getitem__(idx))
    break

